use std::cmp::Ordering;
use std::mem;

/// Enum representing the color of a Red-Black Tree node.
#[derive(Clone, Copy, PartialEq, Debug)]
enum RBTreeNodeColor {
    Red,
    Black,
}

/// Enum representing the side of a Red-Black Tree node.
#[derive(Clone, Copy, PartialEq, Debug)]
enum RBTreeNodeSide {
    Left,
    Right,
}

/// Struct representing a node in the Red-Black Tree.
#[derive(Debug)]
struct RBTreeNode<K, V> {
    color: RBTreeNodeColor,
    key: K,
    value: V,
    parent: Option<Box<RBTreeNode<K, V>>>,
    children: [Option<Box<RBTreeNode<K, V>>>; 2],
}

/// Struct representing the Red-Black Tree.
#[derive(Debug)]
pub struct RBTree<K, V> {
    root_node: Option<Box<RBTreeNode<K, V>>>,
    compare_func: fn(&V, &V) -> Ordering,
    num_nodes: usize,
}

impl<K, V> RBTree<K, V> {
    /// Creates a new Red-Black Tree with the given comparison function.
    pub fn new(compare_func: fn(&V, &V) -> Ordering) -> Self {
        RBTree {
            root_node: None,
            compare_func,
            num_nodes: 0,
        }
    }

    /// Frees the entire subtree rooted at the given node.
    fn free_subtree(&mut self, node: Option<Box<RBTreeNode<K, V>>>) {
        if let Some(mut node) = node {
            self.free_subtree(node.children[0].take());
            self.free_subtree(node.children[1].take());
        }
    }

    /// Frees the entire tree.
    pub fn free(&mut self) {
        self.free_subtree(self.root_node.take());
    }

    /// Returns the side of the node relative to its parent.
    fn node_side(node: &RBTreeNode<K, V>) -> RBTreeNodeSide {
        if let Some(parent) = &node.parent {
            if parent.children[0].as_ref() == Some(node) {
                RBTreeNodeSide::Left
            } else {
                RBTreeNodeSide::Right
            }
        } else {
            RBTreeNodeSide::Left // Default to left if no parent
        }
    }

    /// Returns the sibling of the given node.
    fn node_sibling(node: &RBTreeNode<K, V>) -> Option<Box<RBTreeNode<K, V>>> {
        if let Some(parent) = &node.parent {
            let side = Self::node_side(node);
            parent.children[1 - side as usize].clone()
        } else {
            None
        }
    }

    /// Returns the uncle of the given node.
    fn node_uncle(node: &RBTreeNode<K, V>) -> Option<Box<RBTreeNode<K, V>>> {
        if let Some(parent) = &node.parent {
            Self::node_sibling(parent)
        } else {
            None
        }
    }

    /// Replaces node1 with node2 at its parent.
    fn node_replace(&mut self, node1: &mut RBTreeNode<K, V>, node2: Option<Box<RBTreeNode<K, V>>>) {
        if let Some(node2) = node2 {
            node2.parent = node1.parent.take();
        }

        if node1.parent.is_none() {
            self.root_node = node2;
        } else {
            let side = Self::node_side(node1);
            if let Some(parent) = &mut node1.parent {
                parent.children[side as usize] = node2;
            }
        }
    }

    /// Rotates the subtree rooted at the given node in the specified direction.
    fn rotate(
        &mut self,
        node: &mut RBTreeNode<K, V>,
        direction: RBTreeNodeSide,
    ) -> Option<Box<RBTreeNode<K, V>>> {
        let new_root = node.children[1 - direction as usize].take();

        if let Some(mut new_root) = new_root {
            self.node_replace(node, Some(new_root.clone()));

            node.children[1 - direction as usize] = new_root.children[direction as usize].take();
            new_root.children[direction as usize] = Some(Box::new(node.clone()));

            node.parent = Some(new_root.clone());

            if let Some(child) = &mut node.children[1 - direction as usize] {
                child.parent = Some(Box::new(node.clone()));
            }

            Some(new_root)
        } else {
            None
        }
    }

    /// Inserts a new node with the given key and value into the tree.
    pub fn insert(&mut self, key: K, value: V) -> Option<Box<RBTreeNode<K, V>>> {
        let mut node = Box::new(RBTreeNode {
            color: RBTreeNodeColor::Red,
            key,
            value,
            parent: None,
            children: [None, None],
        });

        let mut parent = None;
        let mut rover = &mut self.root_node;

        while let Some(current) = rover {
            parent = Some(current.clone());

            let side = if (self.compare_func)(&value, &current.value) == Ordering::Less {
                RBTreeNodeSide::Left
            } else {
                RBTreeNodeSide::Right
            };

            rover = &mut current.children[side as usize];
        }

        node.parent = parent;
        *rover = Some(node);

        self.insert_case1(rover.as_mut().unwrap());

        self.num_nodes += 1;

        rover.clone()
    }

    /// Case 1 for insertion: If the new node is at the root of the tree, it must be recolored black.
    fn insert_case1(&mut self, node: &mut RBTreeNode<K, V>) {
        if node.parent.is_none() {
            node.color = RBTreeNodeColor::Black;
        } else {
            self.insert_case2(node);
        }
    }

    /// Case 2 for insertion: If the parent of the new node is red, this conflicts with the Red-Black Tree conditions.
    fn insert_case2(&mut self, node: &mut RBTreeNode<K, V>) {
        if let Some(parent) = &node.parent {
            if parent.color == RBTreeNodeColor::Red {
                self.insert_case3(node);
            }
        }
    }

    /// Case 3 for insertion: If the parent and uncle are both red, repaint them black and the grandparent red.
    fn insert_case3(&mut self, node: &mut RBTreeNode<K, V>) {
        if let Some(uncle) = Self::node_uncle(node) {
            if uncle.color == RBTreeNodeColor::Red {
                if let Some(parent) = &mut node.parent {
                    parent.color = RBTreeNodeColor::Black;
                }
                if let Some(grandparent) = &mut node.parent.as_ref().and_then(|p| p.parent.clone()) {
                    grandparent.color = RBTreeNodeColor::Red;
                }
                if let Some(uncle) = &mut Self::node_uncle(node) {
                    uncle.color = RBTreeNodeColor::Black;
                }
                if let Some(grandparent) = &mut node.parent.as_ref().and_then(|p| p.parent.clone()) {
                    self.insert_case1(grandparent);
                }
            } else {
                self.insert_case4(node);
            }
        } else {
            self.insert_case4(node);
        }
    }

    /// Case 4 for insertion: If the parent is red but the uncle is black, rotate around the parent.
    fn insert_case4(&mut self, node: &mut RBTreeNode<K, V>) {
        let side = Self::node_side(node);
        if side != Self::node_side(&node.parent.as_ref().unwrap()) {
            if let Some(parent) = &mut node.parent {
                self.rotate(parent, 1 - side);
            }
        }
        self.insert_case5(node);
    }

    /// Case 5 for insertion: Rotate at the grandparent and recolor the parent and grandparent.
    fn insert_case5(&mut self, node: &mut RBTreeNode<K, V>) {
        if let Some(parent) = &mut node.parent {
            if let Some(grandparent) = &mut parent.parent {
                let side = Self::node_side(node);
                self.rotate(grandparent, 1 - side);
                parent.color = RBTreeNodeColor::Black;
                grandparent.color = RBTreeNodeColor::Red;
            }
        }
    }

    /// Looks up a node with the given key in the tree.
    pub fn lookup_node(&self, key: K) -> Option<&RBTreeNode<K, V>> {
        let mut node = self.root_node.as_ref();
        while let Some(current) = node {
            match (self.compare_func)(&key, &current.key) {
                Ordering::Equal => return Some(current),
                Ordering::Less => node = current.children[RBTreeNodeSide::Left as usize].as_ref(),
                Ordering::Greater => node = current.children[RBTreeNodeSide::Right as usize].as_ref(),
            }
        }
        None
    }

    /// Looks up the value associated with the given key in the tree.
    pub fn lookup(&self, key: K) -> Option<&V> {
        self.lookup_node(key).map(|node| &node.value)
    }

    /// Returns the number of nodes in the tree.
    pub fn num_entries(&self) -> usize {
        self.num_nodes
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_insert_and_lookup() {
        let mut tree = RBTree::new(|a: &i32, b: &i32| a.cmp(b));
        tree.insert(1, 10);
        tree.insert(2, 20);
        tree.insert(3, 30);

        assert_eq!(tree.lookup(1), Some(&10));
        assert_eq!(tree.lookup(2), Some(&20));
        assert_eq!(tree.lookup(3), Some(&30));
        assert_eq!(tree.lookup(4), None);
    }

    #[test]
    fn test_num_entries() {
        let mut tree = RBTree::new(|a: &i32, b: &i32| a.cmp(b));
        assert_eq!(tree.num_entries(), 0);

        tree.insert(1, 10);
        assert_eq!(tree.num_entries(), 1);

        tree.insert(2, 20);
        assert_eq!(tree.num_entries(), 2);

        tree.insert(3, 30);
        assert_eq!(tree.num_entries(), 3);
    }
}
